import numpy as np
from dataclasses import dataclass
from typing import Literal, Union, Optional
from enum import Enum


class PositionSide(Enum):
    LONG = 1   # 多仓
    SHORT = -1  # 空仓


@dataclass
class Position:
    """
    交易位置信息
    """
    side: PositionSide        # 仓位方向：多/空
    leverage: float           # 杠杆倍数
    open_price: float        # 开仓价格
    size: float              # 仓位大小(USDT)
    wallet_balance: float    # 钱包余额(USDT)
    fee_rate: float = 0.0005  # 开仓手续费率，默认0.05%


class LiquidationCalculator:
    """
    杠杆合约强平计算器
    """

    def __init__(self, maintenance_margin_rate: float = 0.005):
        """
        初始化强平计算器

        Args:
            maintenance_margin_rate: 维持保证金率，默认0.5%
        """
        self.maintenance_margin_rate = maintenance_margin_rate

    def calculate_open_fee(self, position: Position) -> float:
        """
        计算开仓手续费

        Args:
            position: 仓位信息

        Returns:
            float: 开仓手续费
        """
        return position.size * position.fee_rate * 2

    def calculate_initial_margin(self, position: Position) -> float:
        """
        计算初始保证金

        Args:
            position: 仓位信息

        Returns:
            float: 初始保证金金额
        """
        return position.size / position.leverage

    def calculate_contract_qty(self, position: Position) -> float:
        """
        计算合约数量

        Args:
            position: 仓位信息

        Returns:
            float: 合约数量（如BTC数量）
        """
        return position.size / position.open_price

    def calculate_liquidation_price(self, position: Position) -> float:
        """
        计算强平价格

        Args:
            position: 仓位信息

        Returns:
            float: 强平价格
        """
        # 计算初始保证金
        initial_margin = self.calculate_initial_margin(position)

        # 计算开仓手续费
        open_fee = self.calculate_open_fee(position)

        # 实际可用保证金 = 初始保证金 + 钱包余额 - 开仓手续费
        effective_margin = initial_margin + position.wallet_balance - open_fee

        # 计算强平价格
        if position.side == PositionSide.LONG:
            # 多仓强平价格 = 开仓价格 * (1 - (实际可用保证金) / 仓位大小 + 维持保证金率)
            liquidation_price = position.open_price * (
                1 - effective_margin / position.size +
                self.maintenance_margin_rate
            )
        else:
            # 空仓强平价格 = 开仓价格 * (1 + (实际可用保证金) / 仓位大小 - 维持保证金率)
            liquidation_price = position.open_price * (
                1 + effective_margin / position.size -
                self.maintenance_margin_rate
            )

        return liquidation_price

    def calculate_liquidation_risk(self, position: Position) -> dict:
        """
        计算强平风险相关信息

        Args:
            position: 仓位信息

        Returns:
            dict: 包含强平风险相关的详细信息
        """
        initial_margin = self.calculate_initial_margin(position)
        contract_qty = self.calculate_contract_qty(position)
        open_fee = self.calculate_open_fee(position)
        liquidation_price = self.calculate_liquidation_price(position)

        # 计算距离强平的价格变动百分比
        price_change_percentage = abs(liquidation_price - position.open_price) / position.open_price * 100

        # 计算实际可用保证金
        effective_margin = position.wallet_balance - open_fee

        return {
            "side": "多仓" if position.side == PositionSide.LONG else "空仓",
            "leverage": position.leverage,
            "open_price": position.open_price,
            "position_size": position.size-open_fee*position.leverage,
            "wallet_balance": position.wallet_balance,
            "initial_margin": initial_margin,
            "open_fee": open_fee,
            "position_margin": initial_margin-open_fee,
            "effective_margin": effective_margin,
            "contract_quantity": contract_qty,
            "fee_rate": position.fee_rate * 100,  # 转换为百分比
            "maintenance_margin_rate": self.maintenance_margin_rate * 100,  # 转换为百分比
            "liquidation_price": liquidation_price,
            "price_change_to_liquidation": price_change_percentage
        }


def print_liquidation_info(risk_info: dict):
    """
    打印强平风险信息

    Args:
        risk_info: 强平风险信息字典
    """
    print("\n=== 杠杆合约强平风险分析 ===")
    print(f"仓位方向: {risk_info['side']}")
    print(f"杠杆倍数: {risk_info['leverage']}倍")
    print(f"开仓价格: {risk_info['open_price']:.2f} USDT")
    print(f"仓位大小: {risk_info['position_size']:.2f} USDT")
    print(f"开仓手续费率: {risk_info['fee_rate']:.3f}%")
    print(f"初始保证金: {risk_info['initial_margin']:.2f} USDT")
    print(f"开仓手续费: {risk_info['open_fee']:.2f} USDT")
    print(f"仓位保证金: {risk_info['position_margin']:.2f} USDT")
    print(f"实际可用保证金: {risk_info['effective_margin']:.2f} USDT")
    print(f"钱包余额: {risk_info['wallet_balance']:.2f} USDT")
    print(f"合约数量: {risk_info['contract_quantity']:.8f}")
    print(f"维持保证金率: {risk_info['maintenance_margin_rate']:.2f}%")
    print(f"强平价格: {risk_info['liquidation_price']:.2f} USDT")
    print(f"距离强平价格变动: {risk_info['price_change_to_liquidation']:.2f}%")
    print("=" * 30)


if __name__ == "__main__":
    """
    主函数，展示计算器的使用方法
    """
    # 创建计算器实例
    calculator = LiquidationCalculator()

    # 示例：计算多仓强平价格
    long_position = Position(
        side=PositionSide.LONG,
        leverage=100,
        open_price=85000,
        size=311,
        wallet_balance=5,
    )

    # 计算并显示多仓强平风险
    long_risk = calculator.calculate_liquidation_risk(long_position)
    print_liquidation_info(long_risk)

    # 示例：计算空仓强平价格
    short_position = Position(
        side=PositionSide.SHORT,
        leverage=100,
        open_price=83600,
        size=311,
        wallet_balance=4,
    )

    # 计算并显示空仓强平风险
    short_risk = calculator.calculate_liquidation_risk(short_position)
    print_liquidation_info(short_risk)
